/*
 * Wazuh Vulnerability Scanner - Unit Tests
 * Copyright (C) 2015, Wazuh Inc.
 * February 21, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "scanAgentList_test.hpp"
#include "TrampolineOsDataCache.hpp"
#include "TrampolineRemediationDataCache.hpp"
#include "TrampolineSocketDBWrapper.hpp"
#include "scanAgentList.hpp"
#include "shared_modules/utils/mocks/chainOfResponsabilityMock.h"
#include "socketDBWrapperException.hpp"

using testing::_;
using TrampolineScanContext = TScanContext<TrampolineOsDataCache, GlobalData, TrampolineRemediationDataCache>;

const std::string EXPECTED_QUERY_PACKAGE {"agent 001 package get "};
const std::string EXPECTED_QUERY_OS {"agent 001 osinfo get "};
const std::string EXPECTED_QUERY_MANAGER {"agent 0 package get "};
const std::string PACKAGES_RESPONSE {
    R"([{"architecture":"x86_64","checksum":"qwerty","description":"Web Browser","format":"deb","groups":"","install_time":"02/22/2024 00:00:00","item_id":"ytrewq","location":"","multiarch":"","name":"Firefox","priority":"","scan_time":"02/21/2024 00:00:00","size":0,"source":"","vendor":"canonical","version":"122.0.1"},{"architecture":"x86_64","checksum":"asdfgh","description":"Text editor","format":"deb","groups":"","install_time":"02/22/2024 00:00:00","item_id":"hgfdsa","location":"","multiarch":"","name":"Neovim","priority":"","scan_time":"02/21/2024 00:00:00","size":0,"source":"","vendor":"canonical","version":"0.9.5"}])"};

const std::string OS_RESPONSE {
    R"([{"checksum":"qwerty","hostname":"osdata_hostname","os_build":"osdata_build","os_codename":"upstream","os_display_version":"osdata_displayVersion","os_major":"osdata_majorVersion","os_minor":"osdata_minorVersion","os_name":"osdata_name","os_patch":"osdata_patch","os_platform":"osdata_platform","os_release":"osdata_release","os_version":"osdata_version","release":"osdata_release","scan_time":"02/21/2024 00:00:00","sysname":"osdata_sysName","version":"osdata_version"}])"};

const std::string RESPONSE_EMPTY {R"([])"};

TEST_F(ScanAgentListTest, SingleDeleteAndInsertTest)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));
    EXPECT_CALL(*spOsDataCacheMock, setOsData(_, _)).Times(1);

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();
    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(2);
    EXPECT_CALL(*spOsOrchestrationMock, handleRequest(testing::_)).Times(1);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spPackageInsertOrchestrationMock,
                                                                                     spOsOrchestrationMock);

    EXPECT_CALL(*spSocketDBWrapperMock, query(EXPECTED_QUERY_OS, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(nlohmann::json::parse(OS_RESPONSE)));

    EXPECT_CALL(*spSocketDBWrapperMock, query(EXPECTED_QUERY_PACKAGE, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(nlohmann::json::parse(PACKAGES_RESPONSE)));

    nlohmann::json jsonData = nlohmann::json::parse(
        R"({"agent_info":  {"agent_id":"001",  "agent_version":"4.8.0",  "agent_name":"test_agent_name",  "agent_ip":"10.0.0.1",  "node_name":"node01"},  "action":"upgradeAgentDB"})");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);
    contextData->m_agents.push_back({"001", "test_agent_name", "4.8.0", "192.168.0.1"});

    scanAgentList->handleRequest(contextData);
}

TEST_F(ScanAgentListTest, EmptyPackagesWDBResponseTest)
{

    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));
    EXPECT_CALL(*spOsDataCacheMock, setOsData(_, _)).Times(1);

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(0);

    EXPECT_CALL(*spOsOrchestrationMock, handleRequest(testing::_)).Times(1);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spPackageInsertOrchestrationMock,
                                                                                     spOsOrchestrationMock);

    EXPECT_CALL(*spSocketDBWrapperMock, query(EXPECTED_QUERY_OS, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(nlohmann::json::parse(OS_RESPONSE)));

    EXPECT_CALL(*spSocketDBWrapperMock, query(EXPECTED_QUERY_PACKAGE, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(nlohmann::json::parse(RESPONSE_EMPTY)));

    nlohmann::json jsonData = nlohmann::json::parse(
        R"({"agent_info":  {"agent_id":"001",  "agent_version":"4.8.0",  "agent_name":"test_agent_name",  "agent_ip":"10.0.0.1",  "node_name":"node01"},  "action":"upgradeAgentDB"})");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);
    contextData->m_agents.push_back({"001", "test_agent_name", "4.8.0", "192.168.0.1"});

    scanAgentList->handleRequest(contextData);
}

TEST_F(ScanAgentListTest, MutipleRecoverableExceptions)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(2)
        .WillRepeatedly(testing::Throw(SocketDbWrapperException("Temporal error on DB")));

    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(0);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spOsOrchestrationMock,
                                                                                     spPackageInsertOrchestrationMock);

    nlohmann::json jsonData = nlohmann::json::parse(
        R"({"agent_info":  {"agent_id":"001",  "agent_version":"4.8.0",  "agent_name":"test_agent_name",  "agent_ip":"10.0.0.1",  "node_name":"node01"},  "action":"upgradeAgentDB"})");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);
    contextData->m_agents.push_back({"001", "test_agent_name_1", "4.8.0", "192.168.0.1"});
    contextData->m_agents.push_back({"002", "test_agent_name_2", "4.8.0", "192.168.0.2"});

    EXPECT_THROW(scanAgentList->handleRequest(contextData), AgentReScanListException);
    spSocketDBWrapperMock.reset();
}

TEST_F(ScanAgentListTest, OneRecoverableException)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::Throw(SocketDbWrapperException("Temporal error on DB")));

    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(0);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spOsOrchestrationMock,
                                                                                     spPackageInsertOrchestrationMock);

    nlohmann::json jsonData = nlohmann::json::parse(
        R"({"agent_info":  {"agent_id":"001",  "agent_version":"4.8.0",  "agent_name":"test_agent_name",  "agent_ip":"10.0.0.1",  "node_name":"node01"},  "action":"upgradeAgentDB"})");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);
    contextData->m_agents.push_back({"001", "test_agent_name", "4.8.0", "192.168.0.1"});

    EXPECT_THROW(scanAgentList->handleRequest(contextData), AgentReScanListException);
    spSocketDBWrapperMock.reset();
}

TEST_F(ScanAgentListTest, UnrecoverableException)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(2)
        .WillRepeatedly(testing::Throw(std::runtime_error("ERROR on DB")));

    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(0);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spOsOrchestrationMock,
                                                                                     spPackageInsertOrchestrationMock);

    nlohmann::json jsonData = nlohmann::json::parse(
        R"({"agent_info":{"agent_id":"001","agent_version":"4.8.0","agent_name":"test_agent_name","agent_ip":"10.0.0.1","node_name":"node01"},  "action":"upgradeAgentDB"})");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);
    contextData->m_agents.push_back({"001", "test_agent_name", "4.8.0", "192.168.0.1"});

    EXPECT_NO_THROW(scanAgentList->handleRequest(contextData));
    spRemediationDataCacheMock.reset();
    spSocketDBWrapperMock.reset();
    spOsDataCacheMock.reset();
}

TEST_F(ScanAgentListTest, DISABLED_InsertAllTestNotSyncedResponse)
{
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    Os osData {.hostName = "osdata_hostname",
               .architecture = "osdata_architecture",
               .name = "osdata_name",
               .codeName = "upstream",
               .majorVersion = "osdata_majorVersion",
               .minorVersion = "osdata_minorVersion",
               .patch = "osdata_patch",
               .build = "osdata_build",
               .platform = "osdata_platform",
               .version = "osdata_version",
               .release = "osdata_release",
               .displayVersion = "osdata_displayVersion",
               .sysName = "osdata_sysName",
               .kernelVersion = "osdata_kernelVersion",
               .kernelRelease = "osdata_kernelRelease"};

    spOsDataCacheMock = std::make_shared<MockOsDataCache>();
    EXPECT_CALL(*spOsDataCacheMock, getOsData(testing::_)).WillRepeatedly(testing::Return(osData));

    spRemediationDataCacheMock = std::make_shared<MockRemediationDataCache>();
    EXPECT_CALL(*spRemediationDataCacheMock, getRemediationData(testing::_))
        .WillRepeatedly(testing::Return(Remediation {}));

    auto spOsOrchestrationMock = std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    auto spPackageInsertOrchestrationMock =
        std::make_shared<MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>>();

    // Called twice because the server socket response has two packages.
    EXPECT_CALL(*spPackageInsertOrchestrationMock, handleRequest(testing::_)).Times(0);

    auto scanAgentList = std::make_shared<TScanAgentList<TrampolineScanContext,
                                                         MockAbstractHandler<std::shared_ptr<TrampolineScanContext>>,
                                                         TrampolineSocketDBWrapper>>(spOsOrchestrationMock,
                                                                                     spPackageInsertOrchestrationMock);

    nlohmann::json jsonData = nlohmann::json::parse(R"([{"status":"NOT_SYNCED"}])");

    std::variant<const SyscollectorDeltas::Delta*, const SyscollectorSynchronization::SyncMsg*, const nlohmann::json*>
        data = &jsonData;

    auto contextData = std::make_shared<TrampolineScanContext>(data);

    scanAgentList->handleRequest(contextData);
}
